package main

import (
	"github.com/Nik-U/pbc"
	"github.com/consensys/gnark-crypto/ecc/bn254/fr/mimc"
	"math/big"
	rand2 "math/rand"
	"os"
	"sort"
	"time"
)

const storedParams = `type a
    q 8780710799663312522437781984754049815806883199414208211028653399266475630880222957078625179422662221423155858769582317459277713367317481324925129998224791
    h 12016012264891146079388821366740534204802954401251311822919615131047207289359704531102844802183906537786776
    r 730750818665451621361119245571504901405976559617
    exp2 159
    exp1 107
    sign1 1
    sign0 1
`

var rootFloder = "Data"
var fileFloder = "Shen"
var size = "1MB_file.txt"
var testDatafilePath = rootFloder + "/" + fileFloder + "/" + size

var params, _ = pbc.NewParamsFromString(storedParams)
var pairing = params.NewPairing()

var kzg_mimc = mimc.NewMiMC()

func fAdd(a, b *pbc.Element) *pbc.Element {
	ab := pairing.NewZr().Add(a, b)
	return ab
}

func fSub(a, b *pbc.Element) *pbc.Element {
	ab := pairing.NewZr().Sub(a, b)
	return ab
}

func fMul(a, b *pbc.Element) *pbc.Element {
	ab := pairing.NewZr().Mul(a, b)
	return ab
}

func hashFMul(a, b []byte) []byte {
	res := make([]byte, len(a)+len(b))
	var carry uint16
	for i := range a {
		for j := range b {
			prod := uint16(a[i]) * uint16(b[j])
			low := uint8(prod)
			high := uint8(prod >> 8)
			res[i+j] += low
			if res[i+j] < low {
				carry++
			}
			res[i+j+1] += high
			if res[i+j+1] < high {
				carry++
			}
		}
	}
	if carry == 0 {
		return res
	}
	for i := range res {
		res[i] += uint8(carry)
		carry >>= 8
		if carry == 0 {
			break
		}
	}
	return res
}

func fMulZr(a, b *pbc.Element) *pbc.Element {
	return a.Mul(a, b)
}

func fDiv(a, b *pbc.Element) *pbc.Element {
	ab := pairing.NewZr().Div(a, b)
	return ab
}

func fExp(a, b *pbc.Element) *pbc.Element {
	res := a.PowZn(a, b)
	return res
}

func fNeg(a *pbc.Element) *pbc.Element {
	return pairing.NewZr().Neg(a)
}

// polynomialEval evaluates polynomial p at x
func polynomialEval(p []*pbc.Element, x *pbc.Element) *pbc.Element {
	s := pairing.NewZr().Set0()
	for i := 0; i < len(p); i++ {
		xi := pairing.NewZr().PowBig(x, big.NewInt(int64(i)))
		elem := fMul(xi, p[i])
		s = fAdd(s, elem)
	}
	return s
}

func pbcPolynomialiEval(Mx [][]*pbc.Element, x *pbc.Element, i int) *pbc.Element {
	r := pairing.NewZr().Set0()
	for j := range Mx[i] {
		xi := pairing.NewZr().PowBig(x, big.NewInt(int64(j)))
		elem := fMul(xi, Mx[i][j])
		r = fAdd(r, elem)
	}
	return r
}

func hashPbcPolynomialiEval(Mx [][]string, x *pbc.Element, i int) *pbc.Element {
	r := pairing.NewZr().Set0()
	t := pairing.NewZr().Set1()
	for j := range Mx[i] {
		xi := x.PowBig(x, big.NewInt(int64(j)))
		t.SetString(Mx[i][j], 10)
		//println(t)
		elem := pairing.NewZr().Mul(xi, t)
		//println(elem)
		r.Add(r, elem)
	}

	return r
}

func getUniqueRandomNum(sum, choice int) []int {
	selected := make([]int, sum)
	for i := 0; i < sum; i++ {
		selected[i] = i
	}
	rand2.Seed(time.Now().UnixNano())
	rand2.Shuffle(len(selected), func(i, j int) { selected[i], selected[j] = selected[j], selected[i] })
	selected = selected[:choice]
	sort.Ints(selected)

	return selected
}

func polynomialSub(a, b []*pbc.Element) []*pbc.Element {
	maxLen := len(a)
	if len(b) > maxLen {
		maxLen = len(b)
	}
	r := make([]*pbc.Element, maxLen)
	for i := range r {
		r[i] = pairing.NewZr().Set0()
	}
	for i := 0; i < len(a); i++ {
		r[i] = fAdd(r[i], a[i])
	}
	for i := 0; i < len(b); i++ {
		r[i] = fSub(r[i], b[i])
	}
	return r
}

func polynomialMul(a, b []*pbc.Element) []*pbc.Element {
	r := make([]*pbc.Element, len(a)+len(b)-1)
	for i := range r {
		r[i] = pairing.NewZr().Set0()
	}
	for i := 0; i < len(a); i++ {
		for j := 0; j < len(b); j++ {
			temp := fMul(a[i], b[j])
			r[i+j] = fAdd(r[i+j], temp)
		}
	}
	return r
}

func polynomialDiv(a, b []*pbc.Element) ([]*pbc.Element, []*pbc.Element) {
	r := make([]*pbc.Element, len(a)-len(b)+1)
	rem := make([]*pbc.Element, len(a))
	for i := range r {
		r[i] = pairing.NewZr().Set0()
	}
	for i := range rem {
		rem[i] = a[i]
	}

	for len(rem) >= len(b) {
		l := fDiv(rem[len(rem)-1], b[len(b)-1])
		pos := len(rem) - len(b)
		r[pos] = l
		aux := make([]*pbc.Element, pos)
		for i := range aux {
			aux[i] = pairing.NewZr().Set0()
		}
		aux1 := append(aux, l)
		aux2 := polynomialSub(rem, polynomialMul(b, aux1))
		rem = aux2[:len(aux2)-1]
	}
	return r, rem
}

func comparearray(a, b []*pbc.Element) bool {
	if len(a) != len(b) {
		return false
	}
	for i := 0; i < len(a); i++ {
		if a[i].Equals(b[i]) {
			return true
		}
	}
	return false
}

func filePolyGen(blocknum int, tasknum int, segmentnum int) [][]*pbc.Element {
	Mx := make([][]*pbc.Element, blocknum)
	blocksize := segmentnum * 20

	for i := 0; i < blocknum; i++ {
		Mx[i] = make([]*pbc.Element, segmentnum)
		f, _ := os.Open(testDatafilePath)
		f.Seek(int64(tasknum+i*blocksize), 0)
		blockBytes := make([]byte, blocksize)
		f.Read(blockBytes)
		f.Close()
		var x *pbc.Element
		for j := 0; j < segmentnum; j++ {
			mij := blockBytes[20*j : 20*(j+1)]
			if mij != nil {
				x = pairing.NewZr().SetBytes(mij)
			} else {
				x = pairing.NewZr().Set0()
			}
			//fmt.Println(x)
			Mx[i][j] = x
		}
	}
	return Mx
}

func hashFilePolyGen2(blocknum int, tasknum int, segmentnum int) [][][]byte {
	Mx := make([][][]byte, blocknum)
	blocksize := segmentnum * bitNum
	for i := 0; i < blocknum; i++ {
		Mx[i] = make([][]byte, segmentnum)
		f, _ := os.Open(testDatafilePath)
		f.Seek(int64(tasknum+i*blocksize), 0)
		blockBytes := make([]byte, blocksize)
		f.Read(blockBytes)
		f.Close()
		var x []byte
		for j := 0; j < segmentnum; j++ {
			mij := blockBytes[bitNum*j : bitNum*(j+1)]
			if mij != nil {
				//kzg_mimc.Write(mij)
				//x = kzg_mimc.Sum(nil)
				x = mij
			} else {
				x = []byte("")
			}
			Mx[i][j] = x
		}
	}
	return Mx
}
